# pathformer.py
import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.utils.data import Dataset, DataLoader


# 1. Patch Embedding（片段嵌入）
class PatchEmbedding(nn.Module):
    """
    增强稳定性的片段嵌入模块

    输入数据形状：[batch_size, seq_len, input_dim]
    ↓ 经过 PatchEmbedding
    输出形状：  [batch_size, num_patches, embed_dim]
    """


    def __init__(self, input_dim, patch_size, embed_dim):
        super().__init__()
        # ! 将时间序列数据视为一段连续的声音波形
        # ! 分块：把长音频切分成30秒的片段
        self.patch_size = patch_size
        # ! 特征提取：分析每个片段的频谱特征
        # ? nn.Linear()有两个参数，in_features和out_features。
        # ? 宏观上说，这个全连接线性层做的就是将输入进来的特征数量，转换为输出的特征数量
        # ? 其中的W 和 b 的数值训练过程是被 pytorch 省略掉的。
        self.projection = nn.Linear(in_features=input_dim * patch_size, out_features=embed_dim)
        # ! 标准化：统一不同片段的音量
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x):
        batch_size, seq_len, feat_dim = x.size()
        # 动态调整序列长度确保可被patch_size整除
        # ! 这部分是计算 根据当前的序列长度以及分块尺寸，每一段序列能分出多少个块
        num_patches = seq_len // self.patch_size
        # ! valid_len 实际上是对于 seq_len 基础上处理的结果，更加看重可持续性。会将没办法达到分块尺寸的数据给丢弃。
        valid_len = num_patches * self.patch_size
        x = x[:, :valid_len, :] # 自动截断
        # ! 这里就是将矩阵的形状做转换，有点像把多个时间步的数据融合成了一个时间步
        # ? 这里没有搞清楚为什么可以这么做。只知道这里在这么做。
        """
        从原本的：
            # 样本1
            [[1, 2, 3],   # 时间步1
             [4, 5, 6],   # 时间步2
             [7, 8, 9],   # 时间步3
             [10,11,12]], # 时间步4
            
            # 样本2
            [[13,14,15],
             [16,17,18],
             [19,20,21],
             [22,23,24]]
             
        转换到：
            样本1的第一个块 → [1,2,3,4,5,6]  # 时间步1+2
            样本1的第二个块 → [7,8,9,10,11,12] # 时间步3+4
        
        """
        x = x.view(batch_size, num_patches, feat_dim * self.patch_size)
        # ! 这里将调整好数据形状的 x 放到线性层中做维度转换
        x = self.projection(x)
        # ! 转换出来之后进行正则化
        return self.norm(x)


# 2. Positional Encoding（位置编码）
class PositionalEncoding(nn.Module):
    """增强型位置编码"""
    def __init__(self, embed_dim, max_len=5000):
        super().__init__()
        # ! 生成位置索引 [0,1,2,...,max_len-1]
        position = torch.arange(0, max_len).unsqueeze(1)

        # ! 计算频率调节项（指数衰减的除数项）
        # ? 计算div_term = 1/(10000^(2i/d_model))，其中i是维度索引
        # 最后根据embed_dim的长度，出来的是一个一维的数组 (1, embed_dim // 2)
        div_term = torch.exp(torch.arange(0, embed_dim, 2) * (-torch.log(torch.tensor(10000.0)) / embed_dim))
        """
        1. 这里使用了 e 和 log 的关系，为了表达 (10000^(2i/embed_dim)) 这个式子
        # 原理：通过指数衰减生成不同频率的波长（10000^(2i/d_model)的倒数）
        # 效果：低频（大波长）捕捉长程依赖，高频（小波长）捕捉局部模式
        """

        # ! 初始化位置编码矩阵
        # 这部分最后出来的是一个二维数组
        pe = torch.zeros(1, max_len, embed_dim)

        # ! 交替填充正弦和余弦波
        pe[0, :, 0::2] = torch.sin(position * div_term)  # 偶数维度用正弦  0::2是切片的使用
        # 起始位置下标 : 终止位置下标 : 步长，省略代表全部
        pe[0, :, 1::2] = torch.cos(position * div_term)  # 奇数维度用余弦

        # ! 注册为持久化缓冲区（不参与梯度更新）
        self.register_buffer('pe', pe)

    def forward(self, x):

        # + 添加校验逻辑
        if x.size(1) > self.pe.size(1):
            raise ValueError(
                f"输入序列长度{x.size(1)}超过预计算的最大长度{self.pe.size(1)}。"
                f"请初始化时设置更大的max_len参数"
            )
        return x + self.pe[:, :x.size(1), :]


# 3. pathformer 模型（适配多特征预测）
class PathFormer(pl.LightningModule):
    """整合残差连接的时空预测模型"""

    def __init__(self, input_dim=15, patch_size=30, embed_dim=512,
                 num_heads=8, num_layers=6, output_dim=15,
                 dropout=0.2, ff_dim=2048, max_len=5000):
        super().__init__()
        self.save_hyperparameters()

        # 残差连接优化
        # ! 这部分是应对 输入维度 和 输出维度 不匹配做的特殊处理。
        if input_dim != output_dim:
            # ! 这里是让 self.residual 这个变量成为了一个线性变化的容器，可以在之后做维度转换
            # ? 这里的残差是在主路径的处理之外，将最后 序列 的最后一组特征值 加在 主路经处理的结果上。
            """
                ! 主路经 是整体的正常时序预测流程 / 残差路径 是局部的特征值 -直接作用-> 正常预测流程的结果。
            
                # 假设输入是连续7天的气象数据（每天3个特征）
                # x 形状：[32个样本, 7天, 3个特征]
                
                # 取最后一天的全部特征（第7天的温度/湿度/气压）
                residual = x[:, -1, :]  # [32, 3]
                
                # 主路径：分析7天数据中的长期模式
                main_output = transformer_processing(x)  # [32, 15]
                
                # 残差路径：将最后一天的特征转换为预测维度
                residual_output = self.residual(residual)  # [32, 15]
                
                # 综合长期规律与最新观测数据
                final_output = main_output + residual_output  # [32, 15]
            """
            # ? 这种做法可以在一定程度上弥补主路经的判断失误，得到抗过拟合的效果。
            self.residual = nn.Sequential(
                nn.Linear(input_dim, embed_dim),
                nn.Linear(embed_dim, output_dim)
            )
        else:
            # ! 当然如果维度没有不同也就不需要进行维度转换，所以这里就会使 None
            self.residual = None

        # + 修改缩进
        # 核心模块
        # ! 这部分是初始化各个组件，模型的组成相对直接，可以直接对照示意图进行确认。

        # ? 这是片段嵌入模块
        self.patch_embed = PatchEmbedding(input_dim, patch_size, embed_dim) # 时间片段嵌入矩阵
        # ? 这是位置编码模块
        self.pos_encoder = PositionalEncoding(embed_dim, max_len)
        # ? 这是 transformer 编码器模块
        self.transformer = nn.TransformerEncoder(   # Transformer各层的Q/K/V矩阵
            encoder_layer=nn.TransformerEncoderLayer(
                d_model=embed_dim,
                nhead=num_heads,
                dim_feedforward=ff_dim,
                dropout=dropout,
                batch_first=True,
                activation='gelu'
            ),
            num_layers=num_layers
        )
        # ? 这是线性全连接层
        self.fc = nn.Linear(embed_dim, output_dim)
        # ? 这是 dropout 模块
        self.dropout = nn.Dropout(dropout)
        self._init_weights()

    def _init_weights(self):
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_normal_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)

    def forward(self, x):
        # 残差连接优化
        residual = x[:, -1, :] if self.residual else None  # 取最后一个时间步

        x = self.patch_embed(x)
        x = self.pos_encoder(x)
        x = self.transformer(x)
        x = x.mean(dim=1)
        x = self.dropout(x)
        output = self.fc(x)

        if self.residual:
            output += self.residual(residual)

        # ! 这是最后预测数据产生的位置
        return output

    def training_step(self, batch, batch_idx): # <-- 框架自动调用的入口点，应该是pytorch_lightning提供的一个方法，用来训练每个 batch（批次）
        x, y = batch  # 解包数据加载器产生的批次数据
        y_hat = self(x)  # 前向传播（调用forward方法）
        loss = nn.HuberLoss(delta=1.0)(y_hat, y)  # 计算预测值与真实值的损失
        self.log('train_loss', loss, prog_bar=True)  # 记录训练损失
        return loss  # 返回损失供反向传播使用

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            self.parameters(),
            lr=1e-4,
            weight_decay=0.01
        )
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer,
            patience=10,
            mode='min'
        )
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "monitor": "train_loss",
            }
        }


# 4. 数据集类（适配多目标）
class TimeSeriesDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.as_tensor(X, dtype=torch.float32)
        self.y = torch.as_tensor(y, dtype=torch.float32)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]
